from confluent_kafka import Producer
import json
import random
import time
from datetime import datetime
import os
from amap_getter import AMAPDataCollector


def generate_test_data(data=None):

    if data is None:
        data = {
            "timestamp": datetime.now().isoformat() + "Z",
            "location": f"Intersection {random.randint(1, 10)}",
            "vehicle_count": random.randint(0, 20),
            "accident_detected": random.choice([True, False]),
            "traffic_light_status": random.choice(["green", "yellow", "red"])
        }
    return json.dumps(data)


def delivery_report(err, msg):
    if err is not None:
        print(f"Message delivery failed: {err}")
    else:
        print(f"Message delivered to {msg.topic()} [{msg.partition()}]")


class KafkaProducer:
    def __init__(self,
            bootstrap_servers='localhost:9092',
            topic:str='traffic_data',
            location:tuple[int,int]=(121.471873,31.226011)
        ):
        self.producer:Producer = Producer({'bootstrap.servers': bootstrap_servers})
        self.topic:str = topic
        self.dataCollector=AMAPDataCollector(location=location)

    def send_test_data(self, count:int=10, delay:int=1)->None:
        for i in range(count):
            test_data = self.dataCollector.get_data()
            self.producer.produce(
                self.topic, key=str(i),
                value=str(test_data),
                callback=delivery_report
            )
            self.producer.poll(0)
            time.sleep(delay)
        self.producer.flush()


# 使用示例
if __name__ == "__main__":
    kafka_producer = KafkaProducer(bootstrap_servers='localhost:9092', topic='traffic_data')
    kafka_producer.send_test_data(count=10, delay=1)
    os.system("pause")
